import numpy as np
import torch
import torchvision
import torchvision.transforms as transforms
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader, random_split


def load_mnist_data(train_batch_size, test_batch_size, train_ratio):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    # 加载完整的MNIST数据集
    full_trainset = torchvision.datasets.MNIST(root='./datas', train=True, download=True, transform=transform)
    # 计算划分训练集和测试集的样本数量
    train_size = int(train_ratio * len(full_trainset))
    test_size = len(full_trainset) - train_size

    # 划分训练集和测试集
    trainset, testset = random_split(full_trainset, [train_size, test_size])

    # 创建数据加载器
    trainloader = DataLoader(trainset, batch_size=train_batch_size, shuffle=True)
    testloader = DataLoader(testset, batch_size=test_batch_size, shuffle=False)

    return trainloader, testloader


